#ifndef THREAD_POOL_H
#define THREAD_POOL_H
#include <condition_variable>
#include <functional>
#include <future>
#include <queue>
#include <thread>
#include <vector>

class thread_pool
{
    std::queue<std::function<void()>> tasks;
	std::vector<std::thread> workers;
    std::mutex mtx;//For sync usage, protect the `tasks` queue and `stop` flag. 
    std::condition_variable cv;
    bool stop;
public:
	thread_pool(const thread_pool&) = delete;
	thread_pool(thread_pool&&) = delete;
	thread_pool& operator=(const thread_pool&) = delete;
	thread_pool& operator=(thread_pool&&) = delete;

    thread_pool(int thread_count) :stop(false) {

        //create thread_count thread 
        for (size_t i = 0; i < thread_count; ++i) {
            std::thread worker([this]() {
                while (true)
                {
                    /* pop a task from queue, and execute it. */
                    std::function<void()> task;
                    {
                        std::unique_lock<std::mutex> lock(mtx);
                        cv.wait(lock, [this]() { return stop || !tasks.empty(); });
                        if (stop && tasks.empty())
                            return;
                        /* even if stop = 1, once tasks is not empty, then excucte the task until tasks queue become empty*/
                        task = std::move(tasks.front());
                        tasks.pop();
                    }
                    task();
                }
            });
            workers.emplace_back(std::move(worker));
        }
    }
    ~thread_pool()
    {
        //stop thread pool, and notify all threads to finish the remained tasks.
        {
            std::unique_lock<std::mutex> lock(mtx);
            stop = true;
        }
        cv.notify_all();
        for (auto& worker : workers)
            worker.join();
    }

    template <class F, class... Args>
    std::future<std::result_of_t<F(Args...)>> enqueue(F&& f, Args &&...args)
    {
        /* The return type of task `F` */
        using return_type = std::result_of_t<F(Args...)>;

        /* wrapper for no arguments */
        auto task = std::make_shared<std::packaged_task<return_type()>>(
            std::bind(std::forward<F>(f), std::forward<Args>(args)...));

        std::future<return_type> res = task->get_future();
        {
            std::unique_lock<std::mutex> lock(mtx);

            if (stop)
                throw std::runtime_error("The thread pool has been stop.");

            /* wrapper for no returned value */
            tasks.emplace([task]() -> void { (*task)(); });
        }
        cv.notify_one();
        return res;
    }
};

#endif